#!/usr/local/bin/python
import math
import numpy
from numpy.random.mtrand import dirichlet

#@brief Generative process for user state sequence with global influence
#@param n_s : The number of hidden states
#@param n_o : The number of observed states 
#@param alpha_t : The Dirichlet prior of transform matrix from prior state to current state; 
#@param alpha_e : The Dirichlet prior of emission matrix from hidden state to observed value; 
#@param alpha_p : The Dirichlet prior of initial state Pi; 
#@param alpha_s : The Dirichlet prior for selection from friends, global state or historical state;  
#@param alpha_g : The Dirichlet prior for user selection between yesterday's state, global state or friends state 
#@param alpha_l : The Dirichlet prior for friend's influence; 
#@param alpha_a : The Dirichlet prior for interaction behaviour generation; 
#@param T : Lenght of Time.
#@param outfn : The output file name.

def gen_user_state_seq(n_s,
		n_o,
		alpha_t,
		alpha_e,
		alpha_p,
		alpha_s,
		alpha_g,
		alpha_l,
		alpha_a,
		topology,
		T,
		outfn):
	#User has three attribute:
	#	i.	acts:List of actions; 
	#	ii.	phi: Select distribution between himself/herself or global influence;

	# Init global state distribution 	
	theta_g=dirichlet(alpha_g)
	# Init A and B Pi for HMM 
	A=[]		#Transform matrix
	B=[]		#Emission matrix
	Pi=dirichlet(alpha_p)
	for i in range(0,n_s):
		A.append(dirichlet(alpha_t[i]))
		B.append(dirichlet(alpha_e[i]))

	# Init user distribution
	users={}
	for k,v in topology.items():
		d = {}
		d['friends']=v
		d['audience']=retopo[k]
		d['fchoice']=[]
		d['acts']=[]
		d['select']=[]
		d['state']=[]
		d['obs']=[]
		d['phi']=dirichlet(alpha_s)			#Select between himself, global and friends
		d['phi2']=dirichlet(alpha_s[0:2])	#Without any friends
		n_f = len(v)						#number of followers 
		d['zeta']=dirichlet(alpha_a)		#Distribution whether to perform kind of action 
		d['influence']=	dirichlet([alpha_l]*n_f)	#influences from friends
		
		users[k]=d
	
	g=[]
	for t in range(0,T):
		#Genenrate global state
		g.append(gen_multinomial_state(theta_g))
		#Generate user's state and observed value
		for u,k in users.items():
			#Generate user interaction behaviour behaivour 
			if len(k['friends'])>0:		#Only users who has friends can have interactions 
				i_u=gen_multinomial_state(k['influence'])	#Select which friends will influence him
				k['fchoice'].append(k['friends'][i_u])
				a=gen_bi_state(k['zeta'])	#Generate user behaviours
				if 1==a:
					k['acts'].append(k['friends'][i_u])
				else:
					k['acts'].append(0)	#0 means no interaction
			else:
				k['fchoice'].append(0)
				k['acts'].append(0)


			#End generate user behaviour

			if len(k['friends'])>0:
				s=gen_multinomial_state(users[u]['phi'])
			else:
				s=gen_bi_state(users[u]['phi2'])

			users[u]['select'].append(s)
			pre_state=0
			if 0==t:
				pre_state=gen_multinomial_state(Pi)
			else:
				if 0==s:
					pre_state=k['state'][t-1]
				elif 1==s:
					pre_state=g[t]
				elif 2==s:
					#print i_u,' ', k['friends'][i_u]
					pre_state=users[k['friends'][i_u]]['state'][t-1]
				
			cur_s = gen_multinomial_state(A[pre_state])
			users[u]['state'].append(cur_s)
			users[u]['obs'].append(gen_multinomial_state(B[cur_s]))
	
	save_A(outfn,A)
	save_B(outfn,B)
	save_Pi(outfn,Pi)
	save_global_state(outfn,g)
	save_user_state(outfn,users)
	save_user_select(outfn,users)
	save_user_obs(outfn,users)
	save_user_actions(outfn,users)
	save_user_influence(outfn,users)
	save_user_friendschoice(outfn,users)
	save_user_action_zeta(outfn,users)
	save_user_friends(outfn,users)
	save_user_audience(outfn,users)
	return 0
	

def save_A(fn_pre,A):
	fn=open(fn_pre+'_A','w')
	for i in range(0,len(A)):
		for j in range(0,len(A[i])):
			fn.write(str(A[i][j]))
			if j != len(A[i])-1:
				fn.write('\t')
		fn.write('\n')
	fn.close()

def save_B(fn_pre,B):
	fn=open(fn_pre+'_B','w')
	for i in range(0,len(B)):
		for j in range(0,len(B[i])):
			fn.write(str(B[i][j]))
			if j != len(B[i])-1:
				fn.write('\t')
		fn.write('\n')
	fn.close()

def save_Pi(fn_pre,Pi):
	fn=open(fn_pre+'_Pi','w')
	for i in range(0,len(Pi)):
		fn.write(str(Pi[i]))
		if i != len(Pi)-1:
			fn.write('\t')
	fn.close()

def save_global_state(fn_pre,g):
	fn=open(fn_pre+'_global','w')
	for i in range(0,len(g)):
		fn.write(str(g[i]+1))
		if i!=len(g)-1:
			fn.write('\t')
	fn.close()

def save_user_state(fn_pre,users):
	fn=open(fn_pre+'_state','w')
	for u in users.keys():
		for j in range(0,len(users[u]['state'])):
			fn.write(str(users[u]['state'][j]+1))
			if j != len(users[u]['state'])-1:
				fn.write('\t')
		fn.write('\n')
	fn.close()

def save_user_select(fn_pre,users):
	fn=open(fn_pre+'_select','w')
	for u in users.keys():
		for j in range(0,len(users[u]['select'])):
			fn.write(str(users[u]['select'][j]))
			if j != len(users[u]['select'])-1:
				fn.write('\t')
		fn.write('\n')
	fn.close()

def save_user_obs(fn_pre,users):
	fn=open(fn_pre+'_obs','w')
	for u in users.keys():
#		fn.write(str(u))
		for j in range(0,len(users[u]['obs'])):
			fn.write('\t'+str(users[u]['obs'][j]+1))
		fn.write('\n')
	fn.close()

def gen_multinomial_state(dis):
	tmp = numpy.random.multinomial(1,dis,size=1)
	for i in range(0,len(dis)):
		if tmp[0][i]==1:
			return i

def gen_bi_state(dis):
	tmp=numpy.random.binomial(1,dis[1],size=1)
	if tmp[0]==1:
		return 1
	else:
		return 0
def save_user_actions(fn_pre,users):
	fn=open(fn_pre+'_act','w')
	for k,v in users.items():
#		fn.write(str(k))
		for one in v['acts']:
			fn.write(' '+str(one))
		fn.write('\n')
	fn.close()

def save_user_influence(fn_pre,users):
	fn=open(fn_pre+'_influence','w')
	
	for k,v in users.items():
#		fn.write(str(k))
		l = len(v['friends'])
		for i in range(0,l):
			
			fn.write(' '+str(v['influence'][i]))
		fn.write('\n')
	fn.close()
def	save_user_friendschoice(fn_pre,users):
	fn=open(fn_pre+'_fchoice','w')
	for k,v in users.items():
#		fn.write(str(k))
		for one in v['fchoice']:
			fn.write(' '+str(one))
		fn.write('\n')
	fn.close()

def	save_user_action_zeta(fn_pre,users):
	fn=open(fn_pre+'_action_distribution','w')
	for k,v in users.items():
#		fn.write(str(k))
		for one in v['zeta']:
			fn.write(' '+str(one))
		fn.write('\n')
	fn.close()

#@brief Store users's follow list;
def	save_user_friends(fn_pre,users):
	fn=open(fn_pre+'_relation','w')
	for k,v in users.items():
		fn.write(str(k))
		for one in v['friends']:
			fn.write(' '+str(one))
		fn.write('\n')
	fn.close()

#@brief Store users's audience list
def save_user_audience(fn_pre,users):
	fn=open(fn_pre+'_audience','w')
	for k,v in users.items():
		fn.write(str(k))
		for one in v['audience']:
			fn.write(' '+str(one))
		fn.write('\n')
	fn.close()
	



#@brief Get topology from the file. In the topology file, the format is 'ID_A ID_B', means A follows
#		B, and the topology in python would stored in a dict structure, dict={id:[follow_list]}



def get_topology(fn):
	fin =open(fn,'r')
	topo={}		#User and whom he follows
	retopo={}	#User and who follows him
	lines = fin.readlines()
	for one in lines:
		if one.startswith('#'):
			continue

		items=one.split('\t')
		id_a=int(items[0])+1
		id_b=int(items[1])+1
		if not topo.has_key(id_a):
			topo[id_a]=[]
		if not topo.has_key(id_b):
			topo[id_b]=[]
		topo[id_a].append(id_b)

		if not retopo.has_key(id_b):
			retopo[id_b]=[]
		if not retopo.has_key(id_a):
			retopo[id_a]=[]
		retopo[id_b].append(id_a)	
	fin.close()
	return (topo,retopo)
	

if __name__ == '__main__':
	#Configure 1
	'''	
	n_s=4
	n_o=4
	alpha_t=[[100,3,2,1],[2,100,2,1],[1,2,100,2],[1,2,3,100]]
	alpha_e=[[9,1,1,1],[4,7,2,1],[1,4,7,2],[1,2,3,6]]
	alpha_p=[1,1,1,1]
	alpha_s=[8,1,3]		#Prior for pre-state selection
	alpha_g=[1,1,1]
	alpha_l=0.5
	alpha_a=[3,2]
	U=1000
	T=100
	outfn='syns4o4'
	'''
	#Configure 2
	'''
	n_s=3
	n_o=3
	alpha_t=[[20,1,1],[1,15,1],[1,1,10]]
	alpha_e=[[9,1,1],[4,7,2],[2,4,7]]
	alpha_p=[3,2,1]
	alpha_s=[8,1,3]		#Prior for pre-state selection
	alpha_g=[1,1,1]
	alpha_l=0.5
	alpha_a=[3,2]
	T=100
	outfn='synf'
	'''


	#Configure 3
	'''
	n_s=3
	n_o=9
	alpha_t=[[20,1,1],[1,20,1],[1,1,20]]
	alpha_e=[[9,9,9,1,1,1,1,1,1],[1,1,1,10,10,10,1,1,1],[1,1,1,1,1,1,10,10,10]]
	alpha_p=[3,2,1]
	alpha_s=[8,1,3]		#Prior for pre-state selection
	alpha_g=[1,1,1]
	alpha_l=0.5
	alpha_a=[3,2]
	T=100
	outfn='synobs9'
	'''
	
	#configure 4
	'''
	n_s=6
	n_o=3
	alpha_t=[[10,1,1,1,1,1],[1,10,1,1,1,1],[1,1,10,1,1,1],[1,1,1,10,1,1],[1,1,1,1,10,1],[1,1,1,1,1,10]]
	alpha_e=[[9,1,1],[9,1,1],[1,9,1],[1,9,1],[1,1,9],[1,1,9]]
	alpha_p=[3,3,2,2,1,1]
	alpha_s=[8,1,3]		#Prior for pre-state selection
	alpha_g=[1,1,1,1,1,1]
	alpha_l=0.5
	alpha_a=[3,2]
	T=100
	outfn='synstate9'
	'''
	#configure 5	
	n_s=10
	n_o=10
	alpha_t=[]
	alpha_e=[]
	for i in range(0,n_s):
		tmp1=[1]*n_s
		tmp1[i]=50
		alpha_t.append(tmp1)
		tmp2=[1]*n_o
		tmp2[i]=50
		alpha_e.append(tmp2)



	alpha_p=[1]*10;
	alpha_s=[10,1,1]		#Prior for pre-state selection
	alpha_g=[1]*10;
	alpha_l=0.5
	alpha_a=[3,2]
	U=1000
	T=100
	outfn='syns10o10'
	


	(topology,retopo)=get_topology('../../data/graph.txt')
	gen_user_state_seq(n_s,
		n_o,
		alpha_t,
		alpha_e,
		alpha_p,
		alpha_s,
		alpha_g,
		alpha_l,
		alpha_a,
		topology,
		T,
		outfn)

